import os
import sys
import os
import argparse
import logging
from datetime import datetime
import torch
import time

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 导入耕地提取器
from sam2_crop_extractor import extract_crop_fields_from_image
from utils.common_utils import setup_logger
from config import PERFORMANCE_OPTIMIZATION, SAM2_CONFIG, RESULT_FILTER


def main():
    """
    优化版耕地地块提取器启动脚本
    充分利用GPU加速、批处理和异步计算功能
    """
    global PERFORMANCE_OPTIMIZATION
    # 创建参数解析器
    parser = argparse.ArgumentParser(description='优化版遥感耕地地块提取器')
    
    # 必选参数：输入影像路径
    parser.add_argument('-i', '--input', required=True, help='输入遥感影像文件路径（支持TIFF、JPG、PNG等格式）')
    
    # 可选参数
    parser.add_argument('-o', '--output', help='输出矢量文件路径（默认保存在results目录下）')
    parser.add_argument('-f', '--format', choices=['shp', 'geojson'], default='shp', help='输出文件格式（默认为Shapefile）')
    parser.add_argument('-c', '--confidence', type=float, default=SAM2_CONFIG['automatic_mask_generator']['pred_iou_thresh'], 
                      help=f'置信度阈值（默认：{SAM2_CONFIG['automatic_mask_generator']['pred_iou_thresh']}）')
    parser.add_argument('-a', '--area', type=float, default=RESULT_FILTER['min_area'], 
                      help=f'最小耕地面积（平方米，默认：{RESULT_FILTER['min_area']}）')
    parser.add_argument('-l', '--log-level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], 
                      default='INFO', help='日志级别（默认：INFO）')
    
    # 性能优化相关参数
    parser.add_argument('--use-gpu', action='store_true', help='强制使用GPU')
    parser.add_argument('--no-gpu', action='store_true', help='强制不使用GPU')
    parser.add_argument('--batch-size', type=int, help='批处理大小')
    parser.add_argument('--async', action='store_true', dest='use_async', help='启用异步处理')
    parser.add_argument('--no-async', action='store_true', dest='no_async', help='禁用异步处理')
    parser.add_argument('--memory-opt', action='store_true', help='启用内存优化')
    parser.add_argument('--no-memory-opt', action='store_true', help='禁用内存优化')
    
    # 解析参数
    args = parser.parse_args()
    
    # 检查输入文件是否存在
    if not os.path.exists(args.input):
        print(f"错误：输入文件不存在: {args.input}")
        sys.exit(1)
    
    # 创建logs目录
    logs_dir = 'logs'
    os.makedirs(logs_dir, exist_ok=True)
    
    # 设置日志
    log_file = os.path.join(logs_dir, f"optimized_extraction_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log")
    logger = setup_logger(log_file, args.log_level)
    
    # 检查CUDA可用性
    has_cuda = torch.cuda.is_available()
    
    # 根据参数设置性能优化配置
    original_perf_config = PERFORMANCE_OPTIMIZATION.copy()
    
    # 处理GPU设置
    if args.use_gpu:
        if has_cuda:
            PERFORMANCE_OPTIMIZATION['gpu_acceleration']['enabled'] = True
            logger.info("已强制启用GPU加速")
        else:
            logger.warning("CUDA不可用，无法强制使用GPU")
    elif args.no_gpu:
        PERFORMANCE_OPTIMIZATION['gpu_acceleration']['enabled'] = False
        logger.info("已强制禁用GPU加速")
    
    # 处理批处理设置
    if args.batch_size is not None:
        PERFORMANCE_OPTIMIZATION['batch_processing']['batch_size'] = args.batch_size
        logger.info(f"已设置批处理大小: {args.batch_size}")
    
    # 处理异步设置
    if args.use_async:
        PERFORMANCE_OPTIMIZATION['async_processing']['enabled'] = True
        logger.info("已强制启用异步处理")
    elif args.no_async:
        PERFORMANCE_OPTIMIZATION['async_processing']['enabled'] = False
        logger.info("已强制禁用异步处理")
    
    # 处理内存优化设置
    if args.memory_opt:
        PERFORMANCE_OPTIMIZATION['memory_optimization']['enabled'] = True
        logger.info("已强制启用内存优化")
    elif args.no_memory_opt:
        PERFORMANCE_OPTIMIZATION['memory_optimization']['enabled'] = False
        logger.info("已强制禁用内存优化")
    
    # 记录开始信息
    logger.info("=====================================")
    logger.info("      优化版遥感耕地地块提取开始      ")
    logger.info("=====================================")
    logger.info(f"输入文件: {args.input}")
    logger.info(f"输出格式: {args.format.upper()}")
    if args.output:
        logger.info(f"输出文件: {args.output}")
    
    # 记录性能优化配置
    logger.info("\n当前性能优化配置:")
    logger.info(f"- GPU加速: {'已启用' if PERFORMANCE_OPTIMIZATION['gpu_acceleration']['enabled'] else '已禁用'} {f'({torch.cuda.get_device_name(0)})' if has_cuda else ''}")
    logger.info(f"- 批处理: {'已启用' if PERFORMANCE_OPTIMIZATION['batch_processing']['enabled'] else '已禁用'} (批大小: {PERFORMANCE_OPTIMIZATION['batch_processing']['batch_size']})")
    logger.info(f"- 异步处理: {'已启用' if PERFORMANCE_OPTIMIZATION['async_processing']['enabled'] else '已禁用'} (最大线程数: {PERFORMANCE_OPTIMIZATION['async_processing']['max_workers']})")
    logger.info(f"- 内存优化: {'已启用' if PERFORMANCE_OPTIMIZATION['memory_optimization']['enabled'] else '已禁用'}")
    
    # 打印运行信息到控制台
    print("=====================================")
    print("      优化版遥感耕地地块提取器        ")
    print("=====================================")
    print(f"正在处理: {args.input}")
    print(f"输出格式: {args.format.upper()}")
    print(f"日志文件: {log_file}")
    print(f"性能优化: GPU={PERFORMANCE_OPTIMIZATION['gpu_acceleration']['enabled']}, 批处理={PERFORMANCE_OPTIMIZATION['batch_processing']['enabled']}, 异步={PERFORMANCE_OPTIMIZATION['async_processing']['enabled']}")
    print("请稍候，正在提取耕地地块...")
    
    try:
        # 记录开始时间
        start_time = time.time()
        
        # 运行耕地提取器
        output_file = extract_crop_fields_from_image(
            image_path=args.input,
            output_path=args.output
        )
        
        # 计算耗时
        elapsed_time = time.time() - start_time
        
        # 记录和打印完成信息
        logger.info(f"耕地地块提取完成！")
        logger.info(f"结果已保存至: {output_file}")
        logger.info(f"总耗时: {elapsed_time:.2f}秒")
        logger.info("=====================================")
        
        print("\n=====================================")
        print(f"耕地地块提取完成！")
        print(f"结果已保存至: {output_file}")
        print(f"总耗时: {elapsed_time:.2f}秒")
        print("=====================================")
        
        return 0
        
    except Exception as e:
        # 记录错误信息
        logger.error(f"处理过程中发生错误: {e}")
        import traceback
        logger.error(traceback.format_exc())
        
        # 打印错误信息到控制台
        print(f"\n错误: {e}")
        print("请查看日志文件获取详细信息")
        
        return 1
    finally:
        # 恢复原始配置
        PERFORMANCE_OPTIMIZATION = original_perf_config


if __name__ == "__main__":
    sys.exit(main())